#include<cstdlib>
#include<iostream>
#include<cmath>

#include<formic/davidson.h>
#include<formic/random/random.h>
#include<formic/numeric/numeric.h>
#include<formic/numeric/matrix.h>
#include<formic/lapack/interface.h>
//#include<formic/test_driver/utilities.h>

namespace david_test {

  class SimpleMatOp : public formic::MatrixActor<double> {

    private:

      const double * _mat;

    public:

      SimpleMatOp(const double * mat) : _mat(mat) {}

      std::string class_name() { return "david_test::SimpleMatOp"; }

      void operate(const int n, double * const d, double * const y, const int myrank, const int tid) {
        formic::xgemm('N', 'N', n, 1, n, 1.0, _mat, n, d, n, 0.0, y, n);
      }

  };

  class SimplePreconOp : public formic::MatrixActor<double> {

    private:

      const double * _hmat;
      const double * _smat;
      const double & _energy;

    public:

      SimplePreconOp(const double * hmat, const double * smat, const double & energy) : _hmat(hmat), _smat(smat), _energy(energy) {}

      std::string class_name() { return "david_test::SimplePreconOp"; }

      void operate(const int n, double * const d, double * const y, const int myrank, const int tid) {
        for (int i = 0; i < n; i++) {
          double x = _hmat[i+n*i] - _energy * _smat[i+n*i];
          if ( std::abs(x) > 1.0e-6 )
            y[i] = d[i] / x;
          else
            y[i] = 0.0;
        }
      }

  };

}

int main(int argc, char ** argv) {

  formic::mpi::init(argc, argv);

  // open the output file
  formic::open_output_file("output.txt");

  // set random seed
  formic::set_seed(493752);
  //formic::set_seed(std::time(0));

  // set the matrix dimensions
  const int n = 200;

  // create arrays needed by lapack diagonalization routines and matrix inversion function
  std::vector<double> w(n, 0.0);
  std::vector<double> wr(n, 0.0);
  std::vector<double> wi(n, 0.0);
  int lwork = 10 * n;
  std::vector<double> work(lwork, 0.0);
  std::vector<double> rwork(lwork, 0.0);
  int info = 0;
  std::vector<double> A(n*n, 0.0);
  std::vector<int> iwork(2*n, 0);
  double det = 0.0;

  // build the matrices
  std::vector<double> H(n*n, 0.0);
  std::vector<double> S(n*n, 0.0);
  std::vector<double> S_basis(n*n, 0.0);
  for (int j = 0; j < n; j++)
  for (int i = 0; i < n; i++) {
    H[i+n*j] = -10.0 - 2 * formic::random_number<double>();
    //H[i+n*j] += 20 * ( H[i+n*j] > 0.0 ? 1.0 : -1.0 );
    if (i != j) {
      H[i+n*j] = 0.02 * ( formic::random_number<double>() - 0.5 );
      //H[j+n*i] = H[i+n*j];
    }
    S_basis[i+n*j] = 0.1 * ( formic::random_number<double>() - 0.5 );
  }
  formic::xgemm('T', 'N', n, n, n, 1.0, &S_basis.at(0), n, &S_basis.at(0), n, 0.0, &S.at(0), n);
  for (int i = 0; i < n; i++)
    S[i+n*i] += 10.0;
  for (int j = 0; j < n; j++)
  for (int i = 0; i < n; i++)
    S[i+n*j] = S[j+n*i];

  //// set S to the identity
  //for (int j = 0; j < n; j++)
  //for (int i = 0; i < n; i++)
  //  S[i+n*j] = ( i == j ? 1.0 : 0.0 );

  std::vector<double> save_H = H;
  std::vector<double> save_S = S;

//  // transform H by S^(-1/2)
//  std::vector<double> S_neg_half = S;
//  //std::vector<double> S_pos_half = S;
//  {
//
//    // diagonalize S
//    std::vector<double> U = S;
//    formic::xsyev('V', 'U', n, &U.at(0), n, &wr.at(0), &work.at(0), lwork, &rwork.at(0), info);
//    if ( info != 0 )
//      throw formic::Exception("formic::xsyev failed with info = %i") % info;
//
//    // compute S to check
//    A = U;
//    for (int j = 0; j < n; j++)
//    for (int i = 0; i < n; i++)
//      A[i+n*j] *= wr[j];
//    formic::xgemm('N', 'T', n, n, n, 1.0, &U.at(0), n, &A.at(0), n, 0.0, &S_neg_half.at(0), n);
//    for (int j = 0; j < n; j++)
//    for (int i = 0; i < n; i++)
//      if ( std::abs(S[i+n*j] - S_neg_half[i+n*j]) > 1.0e-9 )
//        throw formic::Exception("elment error in S reconstruction of %.2e") % (S[i+n*j] - S_neg_half[i+n*j]);
//
//    // compute S^(-1/2)
//    A = U;
//    for (int j = 0; j < n; j++)
//    for (int i = 0; i < n; i++)
//      A[i+n*j] /= std::sqrt(wr[j]);
//    formic::xgemm('N', 'T', n, n, n, 1.0, &U.at(0), n, &A.at(0), n, 0.0, &S_neg_half.at(0), n);
//
//    // set H to  S^(-1/2) H S^(-1/2)
//    formic::xgemm('N', 'N', n, n, n, 1.0, &S_neg_half.at(0), n, &H.at(0), n, 0.0, &A.at(0), n);
//    formic::xgemm('N', 'N', n, n, n, 1.0, &A.at(0), n, &S_neg_half.at(0), n, 0.0, &H.at(0), n);
//
//    // set S to the identity
//    for (int j = 0; j < n; j++)
//    for (int i = 0; i < n; i++)
//      S[i+n*j] = ( i == j ? 1.0 : 0.0 );
//
//  }

  // prepare the eigenvalue and eigenvector
  double eval = 0.0;
  std::vector<double> evec(n, 0.0);
  for (int i = 0; i < n; i++)
    evec[i] = 2 * ( formic::random_number<double>() - 0.5 );

  // prepare the operators
  david_test::SimpleMatOp Hop(&H.at(0));
  david_test::SimpleMatOp Sop(&S.at(0));
  david_test::SimplePreconOp Mop(&H.at(0), &S.at(0), eval);

  // set the solution threshold and maximum iteration number
  const double thresh = 1.0e-8;
  const int max_it = n;

  // run the davidson solver
  formic::davidson(n, max_it, thresh, eval, &evec.at(0), &Hop, &Sop, &Mop, false);

//  // multiply eigenvecgtor by S^(-1/2)
//  {
//    std::vector<double> temp = evec;
//    formic::xgemm('N', 'N', n, 1, n, 1.0, &S_neg_half.at(0), n, &temp.at(0), n, 0.0, &evec.at(0), n);
//  }

  // check the solution
  {
    H = save_H;
    S = save_S;
    std::vector<double> resid(n, 0.0);
    formic::xgemm('N', 'N', n, 1, n, 1.0, &H.at(0), n, &evec.at(0), n, 0.0, &resid.at(0), n);
    formic::xgemm('N', 'N', n, 1, n, -eval, &S.at(0), n, &evec.at(0), n, 1.0, &resid.at(0), n);
    for (int i = 0; i < n; i++)
      if (std::abs(resid[i]) > thresh)
        throw formic::Exception("large residual element of %.2e") % resid[i];
  }

//  // compute the eigenvalue with lapapck
//  H = save_H;
//  S = save_S;
//  //formic::xsygv(1, 'V', 'U', n, &H.at(0), n, &S.at(0), n, &wr.at(0), &work.at(0), lwork, &rwork.at(0), info);
//  //if ( info != 0 )
//  //  throw formic::Exception("formic::xsygv failed with info = %i") % info;
//  formic::matrix_inverse_lu(n, det, &S.at(0), &A.at(0), &iwork.at(0));
//  formic::xgemm('N', 'N', n, n, n, 1.0, &S.at(0), n, &H.at(0), n, 0.0, &A.at(0), n);
//  formic::xgeev('N', 'N', n, &A.at(0), n, &w.at(0), &wr.at(0), &wi.at(0), (double *)0, n, (double *)0, n, &work.at(0), lwork, &rwork.at(0), info);
//
//  double lapack_eval = wr.at(0);
//  int min_i = 0;
//  for (int i = 0; i < n; i++) {
//    if ( wr.at(i) < lapack_eval ) {
//      lapack_eval = wr.at(i);
//      min_i = i;
//    }
//  }
//  if ( std::abs(wi.at(min_i)) > 1.0e-11 )
//    throw formic::Exception("unexpected imaginary component of %.2e in the lapack eigenvalue") % wi.at(0);
//  if ( std::abs(eval - lapack_eval) > 1.0e-10 )
//    throw formic::Exception("difference of %.2e between davidson and lapack eigenvalues") % ( eval - lapack_eval );

  formic::close_output_file();

  formic::mpi::finalize();

  return EXIT_SUCCESS;

}
